Skip to content

jax comparator #584

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Jul 25, 2025
Merged

jax comparator #584

merged 1 commit into from
Jul 25, 2025

Conversation

@misrasaurabh1
Copy link
Contributor

@misrasaurabh1 misrasaurabh1 commented Jul 25, 2025 •
edited by github-actions bot
Loading

PR Type

Enhancement, Tests

Description

  • Add JAX arrays comparison support

  • Introduce HAS_JAX import flag

  • Extend comparator for jax.Array with allclose

  • Provide comprehensive JAX array tests

Diagram Walkthrough

flowchart LR
  A["comparator(orig, new)"]
  B["HAS_JAX and orig is jax.Array?"]
  C["Check dtype, shape, allclose"]
  D["Fallback to existing logic"]
  A --> B
  B -- "yes" --> C
  B -- "no" --> D
Loading

File Walkthrough
Relevant files
Enhancement
comparator.py
Add JAX array comparison branch                                                   

codeflash/verification/comparator.py

  • Import jax and jax.numpy with HAS_JAX flag
  • Add branch for jax.Array comparison
  • Compare dtype, shape, allclose with NaN support
 +15/-0   
Tests
test_comparator.py
Add comprehensive JAX comparator tests                                     

tests/test_comparator.py

  • Add test_jax with skip if missing JAX
  • Cover scalar, multi-dimensional arrays
  • Test dtype, shape, NaN, infinity, complex, boolean
 +75/-0   

@misrasaurabh1
jax comparator
e15bb84 
Signed-off-by: Saurabh Misra <[email protected]>
@github-actions
Copy link

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪
🧪 PR contains tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

Instance Type Check

When handling JAX arrays, the code only checks orig is jax.Array but does not verify new is also a JAX array, which can lead to attribute errors or incorrect fallthrough behavior. Ensure both operands are JAX arrays before applying JAX-specific comparison logic.

 
if HAS_JAX and isinstance(orig, jax.Array):
    if orig.dtype != new.dtype:
        return False
    if orig.shape != new.shape:
        return False
    return bool(jnp.allclose(orig, new, equal_nan=True))

@github-actions
Copy link

PR Code Suggestions ✨

Explore these optional code suggestions:

CategorySuggestion                                                                                                                                    Impact
Possible issue
Extract boolean with item()

Use .item() to extract a Python boolean scalar from the JAX boolean array result
instead of relying on bool(), which can raise ambiguity errors. This ensures a clean
and reliable boolean return.

codeflash/verification/comparator.py [122]

-return bool(jnp.allclose(orig, new, equal_nan=True))
+return jnp.allclose(orig, new, equal_nan=True).item()
Suggestion importance[1-10]: 7

__

Why: Using .item() converts the JAX boolean array to a Python bool reliably and avoids ambiguous truth-value errors with bool().

Medium
Add new instance type check

Include a check that new is also a JAX array to avoid comparing incompatible types
and potential attribute errors. Mirror the guard used for orig to ensure both inputs
are JAX arrays.

codeflash/verification/comparator.py [117]

-if HAS_JAX and isinstance(orig, jax.Array):
+if HAS_JAX and isinstance(orig, jax.Array) and isinstance(new, jax.Array):
Suggestion importance[1-10]: 6

__

Why: Checking both orig and new for jax.Array prevents attribute errors when comparing non-JAX types and improves type safety.

Low

@misrasaurabh1 misrasaurabh1 enabled auto-merge July 25, 2025 20:18
@misrasaurabh1 misrasaurabh1 merged commit 47e29ec into main Jul 25, 2025
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@KRRT7 KRRT7 KRRT7 approved these changes

Assignees

No one assigned

Labels

Review effort 2/5

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@misrasaurabh1 @KRRT7